"""
@breif 选取掩膜的gui工具
@author Stone at BlueNet, GDOU
@date Nov. 2020
"""
import cv2
import numpy as np


# true and false for mask's element
M_TRUE = 1
M_FALSE = 0

M_COLOR = np.array([0,255,0], dtype=np.uint8)

WINNAME_MASK = 'selected-mask'


def blend(img, mask, color=M_COLOR):
    """Put a color layer on the ``img`` in order to visualize the mask.
    
    known bug: 本来想要实现mask以一定透明度叠加的效果, 但是发现利用当前的函数较难实现.
        因此 blending都在图像副本上做, 以免破坏原图. 目前本函数仅作为可视化工具使用.
        
        正确的blending算法:
        img = img * not(mask) + (img * a + layer * (1-a)) * mask
        其中各项: img * not(mask), mask以外区域, img颜色不变
                  (img * a + layer * (1-a)) * mask, mask 区域内, img与颜色图层layer叠加
    """
    a = 0.5  # opacity
    img = img.copy()
    
    n_channels = max(img.shape[2], len(color))
    
    for i in range(n_channels):
        if color[i] != 0:
            c = int((1-a) * color[i])
            img[:,:,i] = cv2.add(img[:,:,i], c, mask=mask)
    
    return img


def _mouse_callback(event, x, y, flags, params):
    img, mask, radius, winname = params
    print(x, y, flags)  # debug
    
    if flags & cv2.EVENT_FLAG_LBUTTON:
        cv2.circle(mask, (x,y), radius, [M_TRUE], -1)
    elif flags & cv2.EVENT_FLAG_RBUTTON:
        cv2.circle(mask, (x,y), radius, [M_FALSE], -1)
    
    # img = img.copy()  # already done inside blend
    img = blend(img, mask)
    cv2.imshow(winname, img)
    cv2.imshow(WINNAME_MASK, mask*255)


def select_mask(img, radius=10, winname='select_mask'):
    """
    按住左键涂画选区, 按住右键擦除选区, 按Esc返回. 
    注意: 传入的
    
    :param img: 被选取的图像, 假设前三个通道是bgr
    :param winname: 为了避免冲突,可指定串口名
    :param radius: 笔刷的半径
    """
    # img = img.copy()
    mask = np.zeros(img.shape[:2], dtype=np.uint8)
    
    cv2.namedWindow(winname)
    cv2.setMouseCallback(winname, _mouse_callback, (img, mask, radius, winname))
    
    cv2.imshow(winname, blend(img, mask))
    while True:
        if cv2.waitKey(10) & 0xFF == 27: # Esc
            cv2.destroyWindow(winname)
            cv2.destroyWindow(WINNAME_MASK)
            return mask*255


if __name__ == '__main__':
    # im = np.empty([200, 200, 3], dtype=np.uint8)
    im = cv2.imread(r'C:\_AAA\BlueNet\color_seg\a-template.png')
    print(im.dtype)
    
    select_mask(im)

